# -*- encoding: utf-8 -*-
'''
@File    :   mat_reader.py
@Time    :   2021/11/19 10:37
@Author  :   ZhangChaoYang
@Desc    :   读取Matlab文件的工具类
'''

import os
import sys

import numpy as np
from scipy.io import loadmat
from util.signal_processing import fft, sfft, cwt, statistics_feature
import matplotlib.pyplot as plt

class MatFileReader(object):
    '''
    读取Matlab文件的基类
    '''

    def __init__(self, infile, name, trans1dname, trans2dname, work_load=""):
        self.infile = infile
        self.name = name
        self.trans1dname = trans1dname if trans1dname else "original"
        self.trans2dname = trans2dname if trans2dname else "original"
        self.work_load = work_load

    def get_basename(self):
        '''
        获取输入文件的basename
        :return:
        '''
        return os.path.basename(self.infile)

    def read_file_to_dict(self):
        '''
        把整个matlab文件计读到dict里，一列对应一个key
        :return:
        '''
        if not self.infile.endswith(".mat"):
            return {}
        try:
            return loadmat(self.infile)
        except Exception as e:
            print("read matlab file {} failed".format(self.infile))
            print(e)
            return {}

    def window_size(self):
        pass  # 需要子类实现

    def cut_sequence1d(self, sequence, window_size, step_roll_ratio=1.0):
        '''
        把1维长数组切分成很多1维小数组
        :param sequence: 1维长数组
        :param window_size: 每个小数组的长度
        :param step_roll_ratio: 滑动步长占window_size的比例
        :return:
        '''
        if len(sequence) < window_size:
            return None
        move_step = int(window_size * step_roll_ratio)
        if move_step <= 0:
            move_step = window_size
        rect = []
        for begin in range(0, len(sequence) - window_size, move_step):
            end = begin + window_size
            rect.append(sequence[begin:end])
        return np.asarray(rect)

    def cut_sequence2d(self, sequence, window_size, step_roll_ratio=1.0):
        '''
        把2维长数组切分成很多2维小数组
        :param sequence: 2维长数组
        :param window_size: 每个小数组的长度
        :param step_roll_ratio: 滑动步长占window_size的比例
        :return:
        '''
        if len(sequence) < window_size:
            return None
        move_step = int(window_size * step_roll_ratio)
        if move_step <= 0:
            move_step = window_size
        rect = []
        for begin in range(0, len(sequence) - window_size, move_step):
            end = begin + window_size
            rect.append(sequence[begin:end, :])
        return np.asarray(rect)

    def trans1d(self, array, step_roll_ratio, fs):
        '''
        对原始信号进行转换，变换成还是个一维信号。目前支持的转换函数有FFT和统计类特征计算
        '''
        x = self.cut_sequence1d(array, self.window_size(), step_roll_ratio)
        if self.trans1dname == "fft":
            _, x = fft(x, fs)
            return x[:, :-1]
        elif self.trans1dname == "stat":
            x = np.stack(statistics_feature(x), axis=1)
        return x

    def trans2d(self, array, step_roll_ratio, fs):
        '''
        对原始信号进行转换，变换成是个二维信号。目前支持的转换函数有SFFT和CWT
        '''
        x = self.cut_sequence1d(array, self.window_size(), step_roll_ratio)
        if self.trans2dname == "sfft":
            _, _, x = sfft(x, fs)
            x = x[:, :-1, :-1]
        elif self.trans2dname == "cwt":
            results = []
            for sample in x:
                _, matrix = cwt(sample, fs)
                results.append(matrix)
            x = np.asarray(results)
        return x

    def output_normal1d_file(self):
        '''
        :return: 正常样本dunp到哪个文件
        '''
        if self.work_load:
            return os.path.join("corpus", "1d", self.name, self.trans1dname, self.work_load, "normal.npy")
        else:
            return os.path.join("corpus", "1d", self.name, self.trans1dname, "normal.npy")

    def output_anomaly1d_file(self):
        '''
        :return: 异常样本dunp到哪个文件
        '''
        if self.work_load:
            return os.path.join("corpus", "1d", self.name, self.trans1dname, self.work_load, "anomaly.npy")
        else:
            return os.path.join("corpus", "1d", self.name, self.trans1dname, "anomaly.npy")

    def output_normal2d_file(self):
        '''
        :return: 正常样本dunp到哪个文件
        '''
        if self.work_load:
            return os.path.join("corpus", "2d", self.name, self.trans2dname, self.work_load, "normal.npy")
        else:
            return os.path.join("corpus", "2d", self.name, self.trans2dname, "normal.npy")

    def output_anomaly2d_file(self):
        '''
        :return: 异常样本dunp到哪个文件
        '''
        if self.work_load:
            return os.path.join("corpus", "2d", self.name, self.trans2dname, self.work_load, "anomaly.npy")
        else:
            return os.path.join("corpus", "2d", self.name, self.trans2dname, "anomaly.npy")

    def get_label(self):
        '''
        :return: 1正常样本，0异常样本
        '''
        pass  # 需要子类实现

    def read1d(self):
        '''
        从原始数据文件，经过指定的转换函数（变换后还是一维数据），最终变成numpy array
        :return: 可能返回None
        '''
        pass  # 需要子类实现

    def read2d(self):
        '''
        从原始数据文件，经过指定的转换函数（变换后是二维数据），最终变成numpy array
        :return:
        '''
        pass  # 需要子类实现


def readDir1d(reader, dir):
    if not os.path.exists(dir):
        sys.stderr.write("path {} not exists".format(dir))
        return
    x = None
    y = None
    for file in os.listdir(dir):
        infile = os.path.join(dir, file)
        if os.path.isfile(infile):
            reader.infile = infile
            data = reader.read1d()
        elif os.path.isdir(infile):
            data, label = readDir1d(reader, infile)
        if data is not None:
            if reader.get_label() == 1:
                label = np.ones(shape=data.shape[0])
            else:
                label = np.zeros(shape=data.shape[0])
            if x is None:
                x = data
                y = label
            else:
                x = np.vstack([x, data])
                y = np.hstack([y, label])
    return x, y


def readDir2d(reader, dir):
    if not os.path.exists(dir):
        sys.stderr.write("path {} not exists".format(dir))
        return
    x = None
    y = None
    for file in os.listdir(dir):
        infile = os.path.join(dir, file)
        if os.path.isfile(infile):
            reader.infile = infile
            data = reader.read2d()
        elif os.path.isdir(infile):
            data, label = readDir2d(reader, infile)
        if data is not None:
            if reader.get_label() == 1:
                label = np.ones(shape=data.shape[0])
            else:
                label = np.zeros(shape=data.shape[0])
            if x is None:
                x = data
                y = label
            else:
                x = np.vstack([x, data])
                y = np.hstack([y, label])
    return x, y


def convertDir1d(reader, dir_list):
    '''
    对一个目录下的原始数据文件按照指定的转换函数进行转换（变换后还是一维数据），生成一个numpy ndarray，再存入文件
    '''
    normal1d_file = reader.output_normal1d_file()
    anomaly1d_file = reader.output_anomaly1d_file()
    normal1d_dir = os.path.dirname(normal1d_file)
    anomaly1d_dir = os.path.dirname(anomaly1d_file)
    if not os.path.exists(normal1d_dir):
        os.makedirs(normal1d_dir)
    if not os.path.exists(anomaly1d_dir):
        os.makedirs(anomaly1d_dir)

    if not dir_list:
        return
    x, y = None, None
    for dir in dir_list:
        a, b = readDir1d(reader, dir)
        if a is None:
            continue
        if x is None:
            x = a
            y = b
        else:
            x = np.vstack([x, a])
            y = np.hstack([y, b])

    if x is None:
        return
    normal_data = x[y == 1]
    anomaly_data = x[y == 0]
    np.save(normal1d_file, normal_data)
    np.save(anomaly1d_file, anomaly_data)
    print("{} normal data {}".format(reader.name, normal_data.shape[0]))
    print("{} anomaly data {}".format(reader.name, anomaly_data.shape[0]))


def convertDir2d(reader, dir_list):
    '''
    对一个目录下的原始数据文件按照指定的转换函数进行转换（变换后是二维数据），生成一个numpy ndarray，再存入文件
    '''
    normal2d_file = reader.output_normal2d_file()
    anomaly2d_file = reader.output_anomaly2d_file()
    normal2d_dir = os.path.dirname(normal2d_file)
    anomaly2d_dir = os.path.dirname(anomaly2d_file)
    if not os.path.exists(normal2d_dir):
        os.makedirs(normal2d_dir)
    if not os.path.exists(anomaly2d_dir):
        os.makedirs(anomaly2d_dir)

    if not dir_list:
        return
    x, y = None, None
    for dir in dir_list:
        a, b = readDir2d(reader, dir)
        if a is None:
            continue
        if x is None:
            x = a
            y = b
        else:
            x = np.vstack([x, a])
            y = np.hstack([y, b])

    if x is None:
        return
    normal_data = x[y == 1]
    anomaly_data = x[y == 0]
    np.save(normal2d_file, normal_data)
    np.save(anomaly2d_file, anomaly_data)
    print("{} normal data {}".format(reader.name, normal_data.shape[0]))
    print("{} anomaly data {}".format(reader.name, anomaly_data.shape[0]))


def take_a_look_1d(corpus, trans1d):
    '''
    随机抽几个1维样本看一看
    '''
    normal_data_file = os.path.join("corpus", "1d", corpus, trans1d, "normal.npy")
    anomaly_data_file = os.path.join("corpus", "1d", corpus, trans1d, "anomaly.npy")

    X = np.load(normal_data_file)
    ano_X = np.load(anomaly_data_file)
    np.random.shuffle(X)  # 打乱样本顺序
    np.random.shuffle(ano_X)
    N = 6
    N = min(X.shape[0], ano_X.shape[0], N)
    X = X[:N]
    ano_X = ano_X[:N]
    fig, axs = plt.subplots(nrows=2, ncols=N, figsize=(20, 10))
    fig.suptitle(f'{corpus} {trans1d}')
    for i in range(N):
        axs[0][i].set_title(f"正常样本{i}")
        axs[0][i].plot(X[i])
        axs[1][i].set_title(f"异常样本{i}")
        axs[1][i].plot(ano_X[i])
    plt.show()


def take_a_look_2d(corpus, trans2d):
    '''
    随机抽几个2维样本看一看
    '''
    normal_data_file = os.path.join("corpus", "2d", corpus, trans2d, "normal.npy")
    anomaly_data_file = os.path.join("corpus", "2d", corpus, trans2d, "anomaly.npy")

    X = np.load(normal_data_file)
    ano_X = np.load(anomaly_data_file)
    np.random.shuffle(X)  # 打乱样本顺序
    np.random.shuffle(ano_X)
    N = 6
    N = min(X.shape[0], ano_X.shape[0], N)
    X = X[:N]
    ano_X = ano_X[:N]
    fig, axs = plt.subplots(nrows=2, ncols=N, figsize=(20, 10))
    fig.suptitle(f'{corpus} {trans2d}')
    for i in range(N):
        axs[0][i].set_title(f"正常样本{i}")
        axs[0][i].contourf(X[i])
        axs[1][i].set_title(f"异常样本{i}")
        axs[1][i].contourf(ano_X[i])
    plt.show()
